import traceback
from datetime import datetime
import time
import numpy as np
from sentence_transformers import SentenceTransformer
from pymilvus import Collection, CollectionSchema, FieldSchema, DataType, connections, utility
import json


class MilvusClient:
    def __init__(self, host, port):
        try:
            print(f"正在连接 Milvus 数据库，地址：{host}:{port}...")
            connections.connect(alias='default', host=host, port=port)
            print("Milvus连接成功")
        except Exception as e:
            print(f"连接失败: {str(e)}")
            print("详细错误信息：")
            traceback.print_exc()
            raise

    def create_collection(self, table_name):
        if utility.has_collection(table_name):
            return Collection(name=table_name)

        fields = [
            FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
            FieldSchema(name="number", dtype=DataType.VARCHAR, max_length=1),  # 分类编号
            FieldSchema(name="name", dtype=DataType.VARCHAR, max_length=255),  # 分类名称
            FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=768)  # 分类embedding
        ]
        schema = CollectionSchema(fields, description="category collection")
        collection = Collection(name=table_name, schema=schema)
        self.create_index(table_name)
        return collection

    def insert_batch(self, table_name, data):
        collection = self.create_collection(table_name)
        numbers, names, embeddings = zip(*data)  # 解压数据
        insert_result = collection.insert([numbers, names, np.array(embeddings).tolist()])
        collection.compact()
        return insert_result.primary_keys

    def create_index(self, table_name):
        collection = Collection(name=table_name)
        index_params = {"index_type": "IVF_FLAT", "metric_type": "L2", "params": {"nlist": 128}}
        collection.create_index(field_name="embedding", index_params=index_params)

    def delete_all(self, table_name):
        collection = self.create_collection(table_name)
        # 查询所有 id 主键
        result = collection.query("id >= 0")  # 这里只是查询 id 字段，获取所有记录的 id
        # 提取所有 id
        ids = [item['id'] for item in result]  # 如果 result 是列表，直接遍历
        # 执行删除操作，批量删除
        if ids:
            expr = f"id in {ids}"  # 构建删除条件
            collection.delete(expr=expr)  # 执行删除
        collection.compact()

    def query_all(self, table_name):
        collection = self.create_collection(table_name)
        collection.load()
        result = collection.query("id >= 0")
        return len(result)


class ModelUtils:
    _model_instance = None

    @classmethod
    def _load_model(cls, model_path="bert-base-chinese"):
        if cls._model_instance is None:
            cls._model_instance = SentenceTransformer(model_path)
            print("load model finish")

    @classmethod
    def get_embedding(cls, text):
        cls._load_model()
        if text is None:
            raise ValueError("Input text is None")
        embedding = cls._model_instance.encode(text)
        return np.array(embedding)  # 确保返回的是 numpy 数组


def load_category_json(json_file_path):
    with open(json_file_path, "r", encoding="utf-8") as f:
        return json.load(f)


def prepare_data_for_insert(categories):
    data = []
    for number, category in categories.items():
        name = category["name"]
        description = category["description"]
        embedding = ModelUtils.get_embedding(description)
        data.append((number, name, embedding))
    return data


if __name__ == "__main__":
    # Milvus 客户端初始化
    milvus_client = MilvusClient(host="47.106.154.179", port="19530")

    # 加载JSON文件
    categories = load_category_json("categories.json")  # 这里填写你的JSON文件路径

    # 准备插入数据
    data = prepare_data_for_insert(categories)

    table_name = "category"

    len1 = milvus_client.query_all(table_name)
    print("最初的数量为: " + str(len1))

    milvus_client.delete_all(table_name)
    time.sleep(5)
    len2 = milvus_client.query_all(table_name)
    print("删除后的数量为: " + str(len2))

    # 将数据插入Milvus
    milvus_client.insert_batch(table_name=table_name, data=data)
    print("数据插入完成！")

    time.sleep(5)
    len3 = milvus_client.query_all(table_name)
    print("重新插入后的数量为: " + str(len3))


